package io.gatling.extensions.kafka.protocol

import java.util
import java.util.Properties

import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.protocol.{Protocol, ProtocolComponents, ProtocolKey}
import io.gatling.core.session.Session
import io.gatling.core.{CoreComponents, protocol}
import org.apache.kafka.clients.consumer.KafkaConsumer
import org.apache.kafka.common.serialization.StringDeserializer

case class KafkaConsumerProtocol(servers: String, group: String = "test") extends Protocol {

  val props: Properties = new Properties
  props.put("bootstrap.servers", servers)
  props.put("group.id", group)
  props.put("enable.auto.commit", "true")
  props.put("auto.commit.interval.ms", "1000")
  props.put("session.timeout.ms", "30000")
  props.put("auto.offset.reset", "earliest")
  props.put("key.deserializer", classOf[StringDeserializer])
  props.put("value.deserializer", classOf[StringDeserializer])
}

object KafkaConsumerProtocol {

  val kafkaConsumerKey: ProtocolKey[KafkaConsumerProtocol, kafkaConsumerComponents] = new ProtocolKey[KafkaConsumerProtocol, kafkaConsumerComponents] {
    override def protocolClass: Class[protocol.Protocol] = classOf[KafkaConsumerProtocol].asInstanceOf[Class[io.gatling.core.protocol.Protocol]]

    override def defaultProtocolValue(configuration: GatlingConfiguration): KafkaConsumerProtocol = KafkaConsumerProtocol("192.168.106.130:9092,192.168.106.131:9092,192.168.106.132:9092")

    override def newComponents(coreComponents: CoreComponents): KafkaConsumerProtocol => kafkaConsumerComponents = {
      kafkaConsumerProtocol => kafkaConsumerComponents(kafkaConsumerProtocol)
    }
  }
}

case class kafkaConsumerComponents(kafkaConsumerProtocol: KafkaConsumerProtocol) extends ProtocolComponents {

  val consumerMap: util.HashMap[Long, KafkaConsumer[String, String]] = new util.HashMap[Long, KafkaConsumer[String, String]]()

  override def onStart: Session => Session = session => {
    consumerMap.put(session.userId, new KafkaConsumer[String, String](kafkaConsumerProtocol.props))
    session
  }

  override def onExit: Session => Unit = session => {
    consumerMap.get(session.userId).close()
  }
}

case class KafkaConsumerProtocolBuilder(servers: String, group: String = "test") {
  def build: KafkaConsumerProtocol = KafkaConsumerProtocol(servers, group)
}

object KafkaConsumerProtocolBuilder {
  implicit def toKafkaConsumerProtocol(builder: KafkaConsumerProtocolBuilder): KafkaConsumerProtocol = builder.build
}